package ext;

import java.util.*;
import util.*;

public class LeavesShortestPath {
    List<TreeNode> res;
    public List<TreeNode> longestPath(TreeNode root, TreeNode p, TreeNode q) {
        res = new ArrayList<>();
        dfs(root, p, q, new ArrayList<>(), new ArrayList<>());
        return res;
    }

    public TreeNode dfs(TreeNode root, TreeNode p, TreeNode q, List<TreeNode> pList, List<TreeNode> qList) {
        if (root == null || !res.isEmpty()) {
            return null;
        } else if (root == p) {
            pList.add(root);
            return root;
        } else if (root == q) {
            qList.add(root);
            return root;
        }
        TreeNode left = dfs(root.left, p, q, pList, qList);
        TreeNode right = dfs(root.right, p, q, pList, qList);
        if (left == null && right == null) {
            return null;
        } else if (left != null && right != null) {
            for (int i = pList.size() - 1; i >= 0; i--) {
                res.add(pList.get(i));
            }
            res.add(root);
            return null;
        } else {
            TreeNode cur = left == null ? right : left;
            if (cur == p) {
                pList.add(cur);
            } else {
                qList.add(cur);
            }
            return cur;
        }
    }

}
